cdef extern from "stdlib.h":
    ctypedef size_t
    void *malloc(size_t)
    void free(void* ptr)
    int sizeof()
    
cdef extern from "haptree.h":
    ctypedef struct ReadEvidenceList
    
    ctypedef struct HapNode:
        unsigned int nChildren
        HapNode *children
        HapNode *prev_sibling
        HapNode *next_sibling
        HapNode *next
        HapNode *parent
        ReadEvidenceList * evidenceList
        unsigned int allele
        double totalProb
    
    ctypedef struct Locus:
        unsigned int nAlleles
        unsigned int totalNodes
        HapNode *allNodes
        HapNode **alleleIndex
        
    ctypedef struct ReadEvidence:
        unsigned int multiplicity
        double pWeight
        ReadEvidence * next
    
    ctypedef struct HaplotypeCollection:
        unsigned int totalLoci
        unsigned int currentLoci
        Locus **loci
        ReadEvidence evidenceCollection
        unsigned int N

    ctypedef struct PartialRead:
        unsigned int start
        unsigned int length
        unsigned int multiplicity
        unsigned int refCount
        int *alleles

    ctypedef struct Parameters:
        double delta
        double omega
        double significanceCutOff
          
    HaplotypeCollection *createHaplotypeCollection (unsigned int nLoci)
    void destroyHaplotypeCollection (HaplotypeCollection *hapCol)
    void extendHaplotypes(HaplotypeCollection *hapCol, unsigned int nAlleles)
    void printCollection (HaplotypeCollection *hapCol)
    int addRead(HaplotypeCollection *hapCol, PartialRead * read)
    double computeCollectionFreqs(HaplotypeCollection *hapCol, Parameters *pars, double * freqs)
    void setHaplotypeFreqs(HaplotypeCollection *hapCol, double *freqs)
    void pruneHaplotypes(HaplotypeCollection *hapCol, double significance)
    double getListHaplotype(HaplotypeCollection *hapCol, 
                            HapNode *node, unsigned int *haplotype)
    void computeInternalProbs(HaplotypeCollection *hapCol)

import exceptions

class DummyPath(object):
    def __init__(self, id, seq, start):
        self.seq = seq
        self.id = id
        self.start = start
        self.stop = start+len(seq)
    def __str__(self):
        return str(self.seq)

    def __len__(self):
        return len(self.seq)

class DummyIvalMap(object):
    def __init__(self, id, seq, start):
        self.destPath = self.sourcePath = DummyPath(id, seq, start)


cdef class HaplotypeBuilderBase:
    cdef HaplotypeCollection *hapCol
    cdef Parameters params
    def __cinit__(self,unsigned int nLoci, *args, **argv):
        self.hapCol = createHaplotypeCollection(nLoci)
 
    def __dealloc__(self):
        destroyHaplotypeCollection(self.hapCol)

    ## Getter methods
    def printHaplotypes(self):
        printCollection(self.hapCol)

    def __len__(self):
        """ Returns the number of haplotypes in this collection """
        return self.hapCol.loci[self.hapCol.currentLoci-1].totalNodes

    def __iter__(self):
        return iter(self.items())

    def items(self):
        cdef HapNode * node
        cdef unsigned int *haplotype
        
        nLoci = self.hapCol.currentLoci
        node = self.hapCol.loci[nLoci-1].allNodes
        haplotype = <unsigned int *>malloc(sizeof(unsigned int)*nLoci)
        rv = []
        while node != NULL:
            prob = getListHaplotype(self.hapCol, node, haplotype)
            node=node.next
            hap = list()
            for i in xrange(nLoci):
                hap.append(haplotype[i])
            rv.append(  (hap, prob) )
        free(haplotype)
        return rv


    ## Business methods

    def nextLocus(self):
        """Move onto the next locus with nAlleles possible alleles at it.
        Expands the haplotype tree.
        Arguments:
        nAlleles: number of possible alleles at the next locus
        Return value: None
        Exceptions:
        StopIteration: if there are no unexpanded loci.
        """
        if(self.hapCol.currentLoci < self.hapCol.totalLoci):
            extendHaplotypes(self.hapCol, self.nAlleles)
        else:
            raise exceptions.StopIteration("No loci beyond this point")

    def addRead(self, alleleAlphabet, ivalMap, count=1):
        cdef PartialRead read
        cdef int code
        read.start = ivalMap.sourcePath.start
        read.length = len(ivalMap.sourcePath)
        read.multiplicity = count

        read.alleles = <int *> malloc(sizeof(int)*read.length)
        if(read.alleles == NULL):
            raise exceptions.MemoryError("Problems allocating memory for read %s", ivalMap.destPath.id)
        seq=str(ivalMap.destPath)
        try:
            for i in xrange(read.length):
                try:
                    read.alleles[i] = alleleAlphabet[seq[i]]
                except exceptions.KeyError:
                    read.alleles[i] = -1
        except exceptions.KeyError:
            raise exceptions.KeyError("Cannot translate the sequence to allele alphabet. Check the sequence!")
            
        code = addRead(self.hapCol, &read)
        free(read.alleles)
        return code

    def computeFrequencies(self, delta, omega,
                           maxSteps=10000, convergence=1e-6, verbose=False):
        cdef double * freqs
        self.params.delta = delta
        self.params.omega = omega

        freqs = <double *>malloc(sizeof(double)*len(self))
        div=0
        #setHaplotypeFreqs(self.hapCol, NULL)
        for i in xrange(maxSteps):
            if verbose:
                print "Iteration %d:" %(i) 
                self.printHaplotypes()
                print "-------------------"
            div = computeCollectionFreqs(self.hapCol, &self.params, freqs)
            if div<convergence:
                break
        free(freqs)
        computeInternalProbs(self.hapCol)
        return (i+1, div)

    def pruneHaplotypes(self, double threshold):
        pruneHaplotypes(self.hapCol, threshold)

    
class HaplotypeBuilder(HaplotypeBuilderBase):      
    def __init__(self, nLoci, alleleAlphabet='acgt'):
        self.alleleAlphabet=dict(zip(alleleAlphabet.lower()
                                     +alleleAlphabet.upper(),
                                     range(len(alleleAlphabet))*2))
        self.nAlleles = len(alleleAlphabet)
        self.revAlleleAlphabet = dict(zip(range(len(alleleAlphabet)), alleleAlphabet.lower()))

    def addRead(self, ival, count=1):
        return HaplotypeBuilderBase.addRead(self,self.alleleAlphabet, ival, count)
    def items(self):
        rv = []
        for hap,prob in HaplotypeBuilderBase.items(self):
            hapT = ''
            for i in hap:
                hapT+=self.revAlleleAlphabet[i]
            rv.append( (hapT, prob) )
        return rv
